package com.chatmcp.mcprouter.service.impl;

import com.chatmcp.mcprouter.dto.ProxyInfoDto;
import com.chatmcp.mcprouter.dto.ServerConfigDto;
import com.chatmcp.mcprouter.model.SseSession;
import com.chatmcp.mcprouter.service.ProxyService;
import com.chatmcp.mcprouter.service.SessionManager;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

/**
 * Implementation of the SessionManager interface
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class SessionManagerImpl implements SessionManager {

    private final ObjectMapper objectMapper;
    private final ProxyService proxyService;
    
    // Store sessions using ConcurrentHashMap for thread safety (similar to sync.Map in Go)
    private final Map<String, SseSession> sessions = new ConcurrentHashMap<>();
    
    // Store heartbeat schedulers
    private final Map<String, ScheduledExecutorService> heartbeatSchedulers = new ConcurrentHashMap<>();

    @Override
    public SseSession createSession(String key, ServerConfigDto serverConfig, ProxyInfoDto proxyInfo) {
        SseEmitter emitter = new SseEmitter(0L); // No timeout
        SseSession session = new SseSession(emitter, serverConfig, proxyInfo);
        
        // Store session
        storeSession(proxyInfo.getSessionId(), session);
        
        // Store proxy info
        proxyService.storeProxyInfo(proxyInfo.getSessionId(), proxyInfo);
        
        log.debug("Created new session with ID: {}", proxyInfo.getSessionId());
        return session;
    }

    @Override
    public SseSession getSession(String sessionId) {
        SseSession session = sessions.get(sessionId);
        if (session != null && !session.isClosed()) {
            return session;
        }
        
        // Clean up closed sessions
        if (session != null && session.isClosed()) {
            deleteSession(sessionId);
        }
        
        return null;
    }

    @Override
    public void storeSession(String sessionId, SseSession session) {
        sessions.put(sessionId, session);
    }

    @Override
    public void deleteSession(String sessionId) {
        SseSession session = sessions.remove(sessionId);
        if (session != null && !session.isClosed()) {
            session.close();
        }
        
        // Stop heartbeat
        stopHeartbeat(sessionId);
        
        log.debug("Deleted session with ID: {}", sessionId);
    }

    @Override
    public boolean sendEvent(String sessionId, String eventName, Object eventData) {
        SseSession session = getSession(sessionId);
        if (session == null) {
            log.warn("No session found for ID: {}", sessionId);
            return false;
        }
        
        try {
            String data = objectMapper.writeValueAsString(eventData);
            SseEmitter.SseEventBuilder event = SseEmitter.event()
                    .name(eventName)
                    .data(data);
            
            session.getEmitter().send(event);
            return true;
        } catch (IOException e) {
            log.error("Failed to send event to session {}: {}", sessionId, e.getMessage());
            closeSession(sessionId);
            return false;
        }
    }

    @Override
    public boolean sendMessage(String sessionId, String message) {
        SseSession session = getSession(sessionId);
        if (session == null) {
            log.warn("No session found for ID: {}", sessionId);
            return false;
        }
        
        return session.sendMessage(message);
    }

    @Override
    public void startHeartbeat(String sessionId, long intervalMs) {
        // Stop existing heartbeat if any
        stopHeartbeat(sessionId);
        
        // Create new heartbeat scheduler
        ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();
        scheduler.scheduleAtFixedRate(() -> {
            SseSession session = getSession(sessionId);
            if (session == null || session.isClosed()) {
                stopHeartbeat(sessionId);
                return;
            }
            
            try {
                session.getEmitter().send(SseEmitter.event()
                        .name("heartbeat")
                        .data("{\"timestamp\":" + System.currentTimeMillis() + "}"));
                
                log.debug("Sent heartbeat to session: {}", sessionId);
            } catch (IOException e) {
                log.error("Failed to send heartbeat to session {}: {}", sessionId, e.getMessage());
                closeSession(sessionId);
                stopHeartbeat(sessionId);
            }
        }, intervalMs, intervalMs, TimeUnit.MILLISECONDS);
        
        heartbeatSchedulers.put(sessionId, scheduler);
        log.debug("Started heartbeat for session {} with interval {}ms", sessionId, intervalMs);
    }

    @Override
    public void stopHeartbeat(String sessionId) {
        ScheduledExecutorService scheduler = heartbeatSchedulers.remove(sessionId);
        if (scheduler != null) {
            scheduler.shutdownNow();
            log.debug("Stopped heartbeat for session: {}", sessionId);
        }
    }

    @Override
    public void closeSession(String sessionId) {
        SseSession session = getSession(sessionId);
        if (session != null) {
            session.close();
            deleteSession(sessionId);
        }
    }
} 